package anonymouschat.security;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.ObjectInput;
import java.io.ObjectInputStream;
import java.io.ObjectOutput;
import java.io.ObjectOutputStream;
import java.io.Serializable;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.Signature;

import javax.crypto.Cipher;
import javax.crypto.Mac;
import javax.crypto.SecretKey;
import javax.crypto.spec.IvParameterSpec;

import anonymouschat.client.PublicIdentity;
import anonymouschat.exceptions.DecryptionException;
import anonymouschat.messages.securitymessages.AESEncryptedMessage;
import anonymouschat.messages.securitymessages.RSAEncryptedMessage;
import anonymouschat.messages.securitymessages.TextMessageWithSig;
import anonymouschat.security.integrity.DigSig;
import anonymouschat.security.integrity.MacTag;
import anonymouschat.utils.Constants;
import anonymouschat.utils.LocalFileOps;
import anonymouschat.utils.Print;

/**
 * This is the wrapper class of several static security util methods.
 * 
 * AES & RSA; Enc & Dec; Sign & Ver & MAC;
 * 
 * @author bh349
 * 
 */
public class SecurityUtils
{

    /**
     * RSA Encryption for any length object
     * 
     * @param obj
     *            Serializable Object
     * @param pubKey
     *            RSA publicKey
     * @return EncryptedMessage
     */
    public static RSAEncryptedMessage encRSA( Serializable obj,
            PublicKey pubKey, String destUserID, PublicIdentity srcUser ) {
        byte[] result = encRSA( obj, pubKey );
        return new RSAEncryptedMessage( result, destUserID, srcUser );

    }

    /**
     * encrypts an arbitrary-length object
     * 
     * @param obj
     *            the object to be encrypted
     * @param pubKey
     *            they key to encrypt with
     * @return &lt byte[] &gt array of encrypted bytes
     */
    public static byte[] encRSA( Serializable obj, PublicKey pubKey ) {
        byte[] result = null;
        try {
            ByteArrayOutputStream bos = new ByteArrayOutputStream();
            Cipher enc = Cipher.getInstance( Constants.ASYMMETRIC_CYPHER_TYPE );
            enc.init( Cipher.ENCRYPT_MODE, pubKey );
            int offset = 0;
            byte[] unencryptedByteArray = objectToByteArray( obj );
            while( offset < unencryptedByteArray.length ) {

                if( unencryptedByteArray.length - offset < Constants.MAX_ENC_SEGMENT_BYTES ) {
                    byte[] outputBytes = new byte[unencryptedByteArray.length
                            - offset];
                    System.arraycopy( unencryptedByteArray, offset,
                            outputBytes, 0, unencryptedByteArray.length
                                    - offset );
                    bos.write( enc.doFinal( outputBytes ) );
                    break;
                }
                byte[] outputBytes = new byte[Constants.MAX_ENC_SEGMENT_BYTES];
                System.arraycopy( unencryptedByteArray, offset, outputBytes, 0,
                        Constants.MAX_ENC_SEGMENT_BYTES );
                bos.write( enc.doFinal( outputBytes ) );
                offset += Constants.MAX_ENC_SEGMENT_BYTES;
            }
            result = bos.toByteArray();
        } catch( Exception e ) {
            Print.error( "RSA Enc exception", e );
        }
        return result;
    }

    /**
     * AES Encryption
     * 
     * @param obj
     *            Serializable Object
     * @param key
     *            AES SecretKey
     * @param iv
     *            Initialization vector
     * @return EncryptedMessage
     */
    public static AESEncryptedMessage encAES( Serializable obj, SecretKey key,
            byte[] iv ) {
        byte[] result = null;
        try {
            Cipher enc = Cipher.getInstance( Constants.SYMMETRIC_CIPHER_TYPE );
            enc.init( Cipher.ENCRYPT_MODE, key, new IvParameterSpec( iv ) );
            result = enc.doFinal( objectToByteArray( obj ) );
        } catch( Exception e ) {
            Print.error( "AES Enc exception", e );
        }
        return new AESEncryptedMessage( result );
    }

    /**
     * Decrypt RSA
     * 
     * @param obj
     *            EncryptedMessage
     * @param priKey
     *            Private key
     * @throws IOException
     */
    public static TextMessageWithSig decRSA( RSAEncryptedMessage obj,
            PrivateKey priKey ) {
        TextMessageWithSig ret = (TextMessageWithSig) byteArrayToObject( decRSA(
                obj.getEncryptedData(), priKey ) );
        return ret;
    }

    /**
     * RSA Decrypt a byte array to an unencrypted byte array
     * 
     * @param encryptedByteArray
     *            the data to decrypt
     * @param priKey
     *            the key to decrypt with
     * @return
     */
    public static byte[] decRSA( byte[] encryptedByteArray, PrivateKey priKey ) {
        byte[] result = null;
        try {
            Cipher dec = Cipher.getInstance( Constants.ASYMMETRIC_CYPHER_TYPE );
            dec.init( Cipher.DECRYPT_MODE, priKey );

            // Handle arbitrary sized encrypted messages with RSA
            ByteArrayOutputStream bos = new ByteArrayOutputStream();
            int offset = 0;
            int maxDecLength = 384;
            while( offset < encryptedByteArray.length ) {
                byte[] outputBytes;

                if( encryptedByteArray.length - offset < maxDecLength ) {
                    outputBytes = new byte[encryptedByteArray.length - offset];
                    System.arraycopy( encryptedByteArray, offset, outputBytes,
                            0, encryptedByteArray.length - offset );
                    bos.write( dec.doFinal( outputBytes ) );
                    break;
                }

                outputBytes = new byte[maxDecLength];
                System.arraycopy( encryptedByteArray, offset, outputBytes, 0,
                        maxDecLength );
                bos.write( dec.doFinal( outputBytes ) );
                offset += maxDecLength;
            }
            result = bos.toByteArray();
        } catch( Exception e ) {
            LocalFileOps.appendAuditFile( "Unable to decrypt message." );
            Print.error( "Cannot RSA decrypt data", e );
        }
        return result;
    }

    /**
     * Decrypt AES
     * 
     * @param expetedType
     * 
     * @param obj
     *            EncryptedMessage
     * @param key
     *            AES SecretKey
     * @param iv
     *            Initialization vector
     * @throws DecryptionException
     */
    public static Serializable decAES( AESEncryptedMessage obj, SecretKey key,
            byte[] iv ) throws DecryptionException {
        Serializable result = SecurityUtils.byteArrayToObject( decAES(
                obj.getEncryptedData(), key, iv ) );
        return result;
    }

    /**
     * AES Decrypt byte[] to unencrypted byte[]
     * 
     * @param encryptedByteArray
     *            the encrypted data
     * @param key
     *            the secret key
     * @param iv
     *            the initialization vector
     * @return
     * @throws DecryptionException
     */
    public static byte[] decAES( byte[] encryptedByteArray, SecretKey key,
            byte[] iv ) throws DecryptionException {
        Cipher dec = null;
        byte[] ret = null;
        try {
            dec = Cipher.getInstance( Constants.SYMMETRIC_CIPHER_TYPE );
            dec.init( Cipher.DECRYPT_MODE, key, new IvParameterSpec( iv ) );
            ret = dec.doFinal( encryptedByteArray );
        } catch( Exception e ) {
            throw(new DecryptionException(
                    "Error AES decrypting the byte array" ));
        }
        return ret;

    }

    /**
     * Sign
     * 
     * @param obj
     *            Obj to sign
     * @param key
     *            Signing Key
     * @return Signature
     */
    public static DigSig sign( Serializable obj, PrivateKey key ) {
        byte[] signature = null;
        try {
            Signature sign = Signature.getInstance( "SHA256WITHRSA" );
            sign.initSign( key );
            byte[] data = objectToByteArray( obj );
            sign.update( data );
            signature = sign.sign();
        } catch( Exception e ) {
            Print.error( "Digital signature SIGN exception", e );
        }
        return new DigSig( signature );
    }

    /**
     * Ver
     * 
     * @param msg
     *            msg to verify
     * @param sig
     *            Signature
     * @param verKey
     *            Ver key
     * @return true if successfully verified
     */
    public static boolean ver( Serializable msg, DigSig sig, PublicKey verKey ) {
        try {
            Signature ver = Signature.getInstance( "SHA256WITHRSA" );
            ver.initVerify( verKey );
            byte[] data = objectToByteArray( msg );
            ver.update( data );
            return ver.verify( sig.getSig() );
        } catch( Exception e ) {
            Print.error( "Digital signature VER exception", e );
        }
        return false;
    }

    /**
     * key is generated by KeyGenerator.getInstance("HmacSHA256").generateKey()
     * 
     * @param obj
     *            msg to MAC
     * @param key
     *            MAC key
     * @return MACTAG
     */
    public static MacTag mac( Serializable msg, SecretKey key ) {
        byte[] tag = null;
        try {
            Mac m = Mac.getInstance( "HmacSHA256" );
            m.init( key );
            byte[] data = objectToByteArray( msg );
            tag = m.doFinal( data );
        } catch( Exception e ) {
            Print.error( "MAC exception", e );
        }
        return new MacTag( tag );
    }

    /**
     * Convert a byte array to a serializable object
     * @param bytes the bytes representing a serializable object
     * @return &lt Serializable &gt a new object
     */
    public static Serializable byteArrayToObject( byte[] bytes ) {
        if( bytes == null )
            return null;
        ObjectInput out = null;
        Serializable buf = null;
        try {
            // Serialize to a byte array
            ByteArrayInputStream bos = new ByteArrayInputStream( bytes );
            out = new ObjectInputStream( bos );
            buf = (Serializable) out.readObject();

        } catch( Exception e ) {
            Print.error( "Failed to read bytes to object", e );
        } finally {
            if( out != null )
                try {
                    out.close();
                } catch( IOException e ) {

                }
        }
        return buf;
    }

    /**
     * 
     * @param obj
     *            a serializable objec to be converted
     * @return a byte[] array of the object
     * @throws IOException
     */
    public static byte[] objectToByteArray( Serializable obj ) {
        if( obj == null )
            return null;
        ObjectOutput out = null;
        byte[] buf = null;
        try {
            // Serialize to a byte array
            ByteArrayOutputStream bos = new ByteArrayOutputStream();
            out = new ObjectOutputStream( bos );
            out.writeObject( obj );

            // Get the bytes of the serialized object
            buf = bos.toByteArray();

        } catch( IOException e ) {
            Print.error( "Failed to convert obejct to bytes", e );
        } finally {
            if( out != null )
                try {
                    out.close();
                } catch( IOException e ) {

                }
        }
        return buf;
    }

    /**
     * Verify the MAC tag for the given object
     * 
     * @param msg
     *            Object to verify
     * @param tagOfObj
     *            Computed tag to verify
     * @param verKey
     *            MAC key used
     * @return
     */
    public static boolean ver( Serializable msg, MacTag tagOfObj,
            SecretKey verKey ) {
        boolean verified = false;
        try {
            Mac m = Mac.getInstance( "HmacSHA256" );
            m.init( verKey );
            byte[] data = objectToByteArray( msg );
            byte[] t = m.doFinal( data );
            MacTag tag = new MacTag( t );
            if( tag.equals( tagOfObj ) )
                verified = true; // the MAC of the object matched
        } catch( Exception e ) {
            Print.error( "Digital signature VER exception", e );
        }
        return verified;
    }

}
